import os

try:
    import vertexai.preview
    from vertexai.preview.generative_models import GenerativeModel, Content, Part
except:
    pass

from wasabi import msg

from dotenv import load_dotenv

from goldenverba.components.interfaces import Generator

load_dotenv()


class GeminiGenerator(Generator):
    """
    Gemini Generator.
    """

    def __init__(self):
        super().__init__()
        self.name = "Gemini"
        self.description = "Generator using Google's Gemini 1.5 Pro model"
        self.requires_library = ["vertexai"]
        self.requires_env = [
            "GOOGLE_APPLICATION_CREDENTIALS",
            "GOOGLE_CLOUD_PROJECT",
        ]
        self.streamable = True
        self.model_name = os.getenv("GEMINI_MODEL", "gemini-1.5-pro-preview-0409")
        self.context_window = 10000

    async def generate_stream(
        self,
        queries: list[str],
        context: list[str],
        conversation: dict = None,
    ):
        """Generate a stream of response dicts based on a list of queries and list of contexts, and includes conversational context
        @parameter: queries : list[str] - List of queries
        @parameter: context : list[str] - List of contexts
        @parameter: conversation : dict - Conversational context
        @returns Iterator[dict] - Token response generated by the Generator in this format {system:TOKEN, finish_reason:stop or empty}.
        """

        url = os.environ.get("GOOGLE_CLOUD_PROJECT", "")
        if url == "":
            yield {
                "message": "Missing GOOGLE_CLOUD_PROJECT",
                "finish_reason": "stop",
            }

        if conversation is None:
            conversation = {}
        messages = self.prepare_messages(queries, context, conversation)

        try:
            project_id = os.getenv("GOOGLE_CLOUD_PROJECT")

            REGION = "us-central1"
            vertexai.init(project=project_id, location=REGION)

            generative_multimodal_model = GenerativeModel(
                "gemini-1.5-pro-preview-0409",
            )

            completion = await generative_multimodal_model.generate_content_async(
                stream=True, contents=messages
            )

            iter = completion.__aiter__()

            try:
                while True:
                    chunk = await iter.__anext__()
                    if len(chunk.candidates) > 0:
                        if len(chunk.candidates[0].content.parts) > 0:
                            yield {
                                "message": chunk.candidates[0].content.parts[0].text,
                                "finish_reason": chunk.candidates[0].finish_reason,
                            }
                        else:
                            yield {
                                "message": " < Canceled due SAFETY REASONS >",
                                "finish_reason": "",
                            }

            except StopAsyncIteration:
                yield {
                    "message": "",
                    "finish_reason": "stop",
                }
                pass

        except Exception:
            raise

    def prepare_messages(
        self, queries: list[str], context: list[str], conversation: dict[str, str]
    ):
        """
        Prepares a list of messages formatted for a Retrieval Augmented Generation chatbot system, including system instructions, previous conversation, and a new user query with context.

        @parameter queries: A list of strings representing the user queries to be answered.
        @parameter context: A list of strings representing the context information provided for the queries.
        @parameter conversation: A list of previous conversation messages that include the role and content.

        @returns A list of message dictionaries formatted for the chatbot. This includes an initial system message, the previous conversation messages, and the new user query encapsulated with the provided context.

        Each message in the list is a dictionary with 'role' and 'content' keys, where 'role' is either 'system' or 'user', and 'content' contains the relevant text. This will depend on the LLM used.
        """
        messages = []

        for message in conversation:
            messages.append(
                Content(role=message.type, parts=[Part.from_text(message.content)])
            )

        query = " ".join(queries)
        user_context = " ".join(context)

        messages.append(
            Content(
                role="user",
                parts=[
                    Part.from_text(
                        f"{user_context} Please answer this query: '{query}' with this provided context. Only use the context if it is necessary to answer the question."
                    )
                ],
            )
        )

        messages = self.ensure_user_model_alteration(messages)

        return messages

    def ensure_user_model_alteration(self, messages):
        current_role: str = ""

        new_messages: list[Content] = []

        for message in messages:
            if message.role == "system":
                message.role = "model"

        if messages[0].role == "model":
            messages = messages[1:]

        for message in messages:
            if message.role == current_role:
                new_messages[-1] = message
            else:
                new_messages.append(message)
                current_role = message.role

        return new_messages
